import os
import torch.nn as nn
import torch
import torch.nn.functional as F
from torch.nn import Conv2d, Sequential, ReLU
from detection.models.yolo.yolox.yolo_pafpn import YOLOPAFPN
from detection.models.yolo.yolox.yolox import YOLOX, YOLOXHead
from detection.anchor_utils import anchor_utils
from basetrainer.utils import torch_tools


def SeperableConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0):
    """Replace Conv2d with a depthwise Conv2d and Pointwise Conv2d.
    """
    return Sequential(
        Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size,
               groups=in_channels, stride=stride, padding=padding),
        ReLU(True),
        Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1),
    )


def create_model_yolox(net_type,
                       num_classes,
                       width_mult=1.0,
                       pretrained=True,
                       is_test=False,
                       device="cuda:0",
                       **kwargs):
    """
    https://github.com/Megvii-BaseDetection/YOLOX
    yolox-nano, yolox-tiny, yolox-s, yolox-m(64), yolox-l, yolox-x
    :param net_type: {"yolox-nano","yolox-tiny","yolox-s","yolox-m","yolox-l","yolox-x",}
    :param depthwise: True or False
    :param act: {"silu","relu","lrelu"}
    """

    def init_yolo(M):
        for m in M.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eps = 1e-3
                m.momentum = 0.03

    param_dict = {
        "yolox-nano": (0.33, 0.25),
        "yolox-tiny": (0.33, 0.375),
        "yolox-s": (0.33, 0.50),
        "yolox-m": (0.67, 0.75),
        "yolox-l": (1.0, 1.0),
        "yolox-x": (1.33, 1.25),
    }
    if not net_type in param_dict:
        raise Exception(" YOLOX only support net_type={},not support {}".format(list(param_dict.keys()), net_type))
    depthwise = kwargs["depthwise"] if "depthwise" in kwargs else False
    # act={"silu","relu","lrelu"}
    act = kwargs["act"] if "act" in kwargs else "silu"
    depth, width = param_dict[net_type]
    backbone = YOLOPAFPN(depth, width, width_mult=width_mult, depthwise=depthwise, act=act)
    depth, width = param_dict[net_type]
    in_channels = [256, 512, 1024]
    # backbone = YOLOPAFPN(depth, width, in_channels=in_channels)
    head = YOLOXHead(num_classes, width, in_channels=in_channels, is_test=is_test, depthwise=depthwise, act=act)
    model = YOLOX(backbone, head, is_test=is_test)
    model.apply(init_yolo)
    model.head.initialize_biases(1e-2)
    if pretrained:
        file = pretrained if isinstance(pretrained, str) else get_pretrained_url(net_type)
        if os.path.exists(file):
            print("use pretrained file {}".format(file))
            state_dict = torch.load(file, map_location=device)
            # load_dict = torch.load(file, map_location=device)["model"]
            if "model" in state_dict:
                state_dict = state_dict["model"]  # 原始官方模型：yolox_nano_h.pth
            model = torch_tools.load_pretrained_model(model, state_dict)
            # 重新保存模型，以兼容低版本torch
            # torch.save(model.state_dict(), file+".low", _use_new_zipfile_serialization=False)
        else:
            raise Exception("no pretrained file".format(file))
    return model


def get_pretrained_url(mode='small'):
    """
    :param mode:
    :param width_mult:
    :return:
    """
    if mode == 'yolox-nano':
        file = "yolox_nano.pth"
    elif mode == 'yolox-tiny':
        file = "yolox_tiny.pth"
    elif mode == 'yolox-s':
        file = "yolox_s.pth"
    else:
        return ""
    file = os.path.join(os.path.dirname(__file__), "pretrained", file)
    return file
